package cim144

import chisel3._
import chisel3.util._
import scala.io.Source

trait Mvm_config{
  val ROW_NUM = 576
  val COL_NUM = 128
  val ROW_BW  = log2Ceil(ROW_NUM+1)
  val COL_BW  = log2Ceil(COL_NUM+1)
  //  val SAVE_WIDTH = 16
}

class cim_rom extends BlackBox with Mvm_config {
  val io = IO(new Bundle{
    val a   = Input(UInt(ROW_BW.W))
    val spo = Output(UInt((COL_NUM*4).W))
  })
}

class cim_rom_sim extends Module with Mvm_config {
  val io = IO(new Bundle{
    val a   = Input(UInt(ROW_BW.W))
    val spo = Output(UInt((COL_NUM*4).W))
  })
  //val mem = VecInit("h12".U,"h81".U,"hd1".U,"hff".U)
  val mem = VecInit(Seq.fill(ROW_NUM)(0.U((COL_NUM*4).W)))
  for(i<-0 until ROW_NUM) {
    val tmp = VecInit((0 until COL_NUM).map(t => 1.asUInt(4.W)))
    if(i < ROW_NUM-3){
      tmp(COL_NUM-1) := 2.U
    }else{
      tmp(0) := 2.U
    }
    mem(i) := tmp.asUInt
  }
  io.spo := mem(io.a)
}

class cim_rom_sim_file extends Module with Mvm_config {
  val io = IO(new Bundle{
    val a   = Input(UInt(ROW_BW.W))
    val spo = Output(UInt((COL_NUM*4).W))
  })
  val mem = VecInit(Seq.fill(ROW_NUM)(0.U((COL_NUM*4).W)))
  val source = Source.fromFile("/home/hezf/chipyard_docker/generators/cim144/mlp_weight_rom.txt")
  val lines  = source.getLines().map(x=> x.split(',')).toList

  for(i<-0 until ROW_NUM) {
    val tmp1 = lines(i).map(x=> x.toInt.asUInt(4.W))
    val tmp2 = VecInit((0 until COL_NUM).map(t => tmp1(COL_NUM-1-t)))
    mem(i) := tmp2.asUInt
  }
  io.spo := mem(io.a)
}


class mvm_IO extends Bundle with Mvm_config {
  val row_begin = UInt(ROW_BW.W)
  val col_begin = UInt(COL_BW.W)
  val row_end = UInt(ROW_BW.W)
  val col_end = UInt(COL_BW.W)
}

class cim_mvm_8bit_shift(save_width : Int ) extends Module with Mvm_config {
  val io = IO(new Bundle{
    val start     = Input(Bool())
    val mvm_done  = Output(Bool())
    val rcbd      = Input(new mvm_IO)
    val push_buf  = Input(UInt((ROW_NUM*2).W))
    val save_buf  = Output(UInt((COL_NUM*save_width).W))
    val shift_num = Input(UInt(8.W))
  })
  val idle :: clr::mvm :: Nil = Enum(3)
  val state = RegInit(idle)
  switch(state){
    is(idle){   state := Mux(io.start,clr,idle)    }
    is(clr) {   state := mvm   }
    is(mvm) {   state := Mux(io.mvm_done,idle,mvm)}
  }

  val input_buf = io.push_buf.asTypeOf(Vec(ROW_NUM,UInt(2.W)))
  val output_buf = VecInit(Seq.fill(COL_NUM)(0.S((save_width).W)))
  val output_buf_tmp = RegInit(VecInit(Seq.fill(COL_NUM)(0.S((save_width*2).W))))
  val output_buf_tmp_shift = VecInit(Seq.fill(COL_NUM)(0.S((save_width*2).W)))
  val addr = RegInit(UInt(ROW_BW.W),0.U)
  when(state === clr){
    addr := io.rcbd.row_begin
  }.elsewhen(io.mvm_done){   //1.avoid out of range
    addr := 0.U
  }.elsewhen(state===mvm){
    addr := addr + 1.U
  }
  //rom
  // if fpga
  //val rom = Module(new cim_rom())
  // if simulation  
  //val rom = Module(new cim_rom_sim())
  // if simulation ,load weight from romfile
  val rom = Module(new cim_rom_sim_file())

  rom.io.a := Mux(addr<ROW_NUM.U,addr,(ROW_NUM-1).U)
  val rom_out_reverse = (rom.io.spo).asTypeOf(Vec(COL_NUM,SInt(4.W)))
  val rom_out = rom_out_reverse.reverse
  val output_en = Wire(Vec(COL_NUM,Bool()))
  val add_num   = Wire(Vec(COL_NUM,SInt((save_width*2).W)))
  for (i <-0 until COL_NUM ){
    output_en(i) := ( i.U >=io.rcbd.col_begin)& (i.U <= io.rcbd.col_end)
    add_num(i) := rom_out(i)

    when(state === clr){
      output_buf_tmp(i) := 0.S
    }.elsewhen(state===mvm){
      output_buf_tmp(i) := Mux(output_en(i), output_buf_tmp(i) +
          Mux(input_buf(addr)(1).asBool(), -add_num(i),
          Mux(input_buf(addr)(0).asBool(),  add_num(i), 0.S)), 0.S)
    }.elsewhen(state===idle){
      output_buf_tmp(i) := output_buf_tmp(i)
    }
    output_buf_tmp_shift(i):= output_buf_tmp(i) >> io.shift_num
    output_buf(i) := Mux(output_buf_tmp_shift(i)>=127.S,127.S,Mux(output_buf_tmp_shift(i)<= -127.S,-127.S,output_buf_tmp_shift(i)))
  }
  //val output_buf_reverse = Wire(Vec(COL_NUM,SInt(16.W)))  //2.weight is BigEndian
  //output_buf_reverse := output_buf.reverse
  //io.save_buf := output_buf_reverse.asUInt()
  io.save_buf := output_buf.asUInt()
  io.mvm_done := (addr === io.rcbd.row_end)&(state===mvm)//3.state ==mvm
}

object cim_mvm_8bit_shift{
  def apply(start:Bool,rcbd:mvm_IO,push_buf:UInt,shift_num:UInt):(Bool,UInt)={
    val inst = Module(new cim_mvm_8bit_shift(8))
    inst.io.start := start
    inst.io.rcbd  := rcbd
    inst.io.push_buf := push_buf
    inst.io.shift_num:= shift_num
    (inst.io.mvm_done,inst.io.save_buf)
  }
}